"""
自我一致性
"""
from dotenv import load_dotenv
from openai import OpenAI

# 加载环境变量
load_dotenv()

# 创建 openAI 客户端
client = OpenAI()

# 定义一个问题，询问当自己 70 岁，6 岁时妹妹是自己年龄一半，现在妹妹的年龄
prompt2 = """
现在我70岁了，当我6岁时，我的妹妹是我的年龄的一半。现在我的妹妹多大？请逐步思考
"""
# 定义一系列问答示例，用于给模型提供解题思路和格式，最后提出与 prompt2 相同的问题
prompt = """
Q：林中有15棵树。林业工人今天将在林中种树。完成后，将有21棵树。林业工人今天种了多少棵树？
A：我们从15棵树开始。后来我们有21棵树。差异必须是他们种树的数量。因此，他们必须种了21 - 15 = 6棵树。答案是6。

Q：停车场有3辆汽车，又来了2辆汽车，停车场有多少辆汽车？
A：停车场已经有3辆汽车。又来了2辆。现在有3 + 2 = 5辆汽车。答案是5。

Q：Leah有32块巧克力，她的姐姐有42块。如果他们吃了35块，他们总共还剩多少块？
A：Leah有32块巧克力，Leah的姐姐有42块。这意味着最初有32 + 42 = 74块巧克力。已经吃了35块。因此，他们总共还剩74 - 35 = 39块巧克力。答案是39。

Q：Jason有20个棒棒糖。他给Denny一些棒棒糖。现在Jason只有12个棒棒糖。Jason给Denny多少棒棒糖？
A：Jason有20个棒棒糖。因为他现在只有12个，所以他必须把剩下的给Denny。他给Denny的棒棒糖数量必须是20 - 12 = 8个棒棒糖。答案是8。

Q：Shawn有五个玩具。圣诞节，他从他的父母那里得到了两个玩具。他现在有多少个玩具？
A：他有5个玩具。他从父母那里得到了2个，所以在那之后他有5 + 2 = 7个玩具。答案是7。

Q：服务器房间里有9台计算机。从周一到周四，每天都会安装5台计算机。现在服务器房间里有多少台计算机？
A：从周一到周四有4天。每天都添加了5台计算机。这意味着总共添加了4 * 5 = 20台计算机。一开始有9台计算机，所以现在有9 + 20 = 29台计算机。答案是29。

Q：Michael有58个高尔夫球。星期二，他丢失了23个高尔夫球。星期三，他又丢失了2个。星期三结束时他还剩多少个高尔夫球？
A：Michael最初有58个球。星期二他丢失了23个，所以在那之后他有58 - 23 = 35个球。星期三他又丢失了2个，所以现在他有35 - 2 = 33个球。答案是33。

Q：Olivia有23美元。她用每个3美元的价格买了五个百吉饼。她还剩多少钱？
A：她用每个3美元的价格买了5个百吉饼。这意味着她花了3 * 5 = 15美元。她还剩23 - 15 = 8美元。答案是8。

Q：现在我70岁了，当我6岁时，我的妹妹是我的一半年龄。现在我的妹妹多大？
A：
"""


# 定义一个函数 用于生成多个不同的答案
def get_multiple_completions(prompt, n=5, model="qwen-max-2025-01-25", temperature=0.7):
    # 用于存储答案
    responses = []
    # 循环 n 次，每次调用 OpenAI Api 生成一个答案
    for _ in range(n):
        response = client.chat.completions.create(
            model=model,
            messages=[{"role": "user", "content": prompt}],
            temperature=temperature
        )
        # 从响应中提取答案 并去掉首位空格
        answer = response.choices[0].message.content.strip()
        # 将答案添加到列表中
        responses.append(answer)

    # 返回所有答案
    return responses


# 定义一个函数，用来大模型来统计和选择一致性的答案
def get_most_common_answer_via_model(responses, model="qwen-max-2025-01-25"):
    # 将所有答案用换行符连成一个字符串
    responses_str = "\n".join(responses)
    # 构造新的提示词，让模型从这些答案中选择出现最多的作为最终答案
    prompt_for_model = f"""请根据以下生成的多个答案，选择出现次数最多的答案作为最终答案：{responses_str}最终答案是："""
    response = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt_for_model}],
        temperature=0
    )
    # 从响应中提取最终答案 并去掉首位空格
    final_answer = response.choices[0].message.content.strip()
    # 返回最终答案
    return final_answer


# 调度函数
def get_completion_with_self_consistency(prompt, n=5):
    responses = get_multiple_completions(prompt, n)
    print("Generated Responses:")
    for i, response in enumerate(responses):
        print(f"Response {i + 1}: {response}")
        # 调用大模型来选择一致性的答案
    final_answer = get_most_common_answer_via_model(responses)
    print(f"Final Answer: {final_answer}")


print("Final Answer for prompt:")
# get_completion_with_self_consistency(prompt)


print(f"aaa :{prompt}")
